# 取交集result1.json和result2.json的count_list，并保存到result3.json
# result1.json是一个list，result2.json是一个dict，result3.json也是一个list
# result*.json 路径是{project_root}/data/result*.json
# result2.json的股票做交集时只取每个股票的前6个字符
# 使用loguru

import json
import os
from loguru import logger

def main():
    # 获取项目根目录
    project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
    
    # 构建文件路径
    result1_path = os.path.join(project_root, "data", "result1.json")
    result2_path = os.path.join(project_root, "data", "result2.json")
    result3_path = os.path.join(project_root, "data", "result3.json")
    
    logger.info(f"读取文件: {result1_path}")
    logger.info(f"读取文件: {result2_path}")
    
    try:
        # 读取result1.json
        with open(result1_path, 'r', encoding='utf-8') as f:
            result1_data = json.load(f)
        logger.info(f"result1.json 包含 {len(result1_data)} 个股票代码")
        
        # 读取result2.json
        with open(result2_path, 'r', encoding='utf-8') as f:
            result2_data = json.load(f)
        logger.info(f"result2.json 包含 {len(result2_data['count_list'])} 个股票代码")
        
        # 取交集
        set1 = set(result1_data)
        # result2.json的股票做交集时只取每个股票的前6个字符
        set2 = set([code[:6] for code in result2_data['count_list']])
        intersection = set1.intersection(set2)
        
        # 转换为列表并排序
        result3_data = sorted(list(intersection))
        
        logger.info(f"交集包含 {len(result3_data)} 个股票代码")
        logger.info(f"交集股票代码: {result3_data}")
        
        # 保存到result3.json
        with open(result3_path, 'w', encoding='utf-8') as f:
            json.dump(result3_data, f, ensure_ascii=False, indent=2)
        
        logger.success(f"结果已保存到: {result3_path}")
        
    except FileNotFoundError as e:
        logger.error(f"文件未找到: {e}")
    except json.JSONDecodeError as e:
        logger.error(f"JSON解析错误: {e}")
    except Exception as e:
        logger.error(f"发生错误: {e}")

if __name__ == "__main__":
    main()
